import nibabel as nib
import numpy as np
import SimpleITK as sitk
import os
import shutil

def adjustMethod1(data_resampled,w_width,w_center):
    val_min = w_center - (w_width / 2)
    val_max = w_center + (w_width / 2)
    print(val_max,val_min)

    data_adjusted = data_resampled.copy()
    data_adjusted[data_resampled < val_min] = val_min
    data_adjusted[data_resampled > val_max] = val_max

    return data_adjusted

files = os.listdir('../autodl-tmp/point_train_data/')
for file in files:
    print(file)
    
#     seg_path = os.path.join('../autodl-tmp/point_train_data/',file,'segmentation.nii.gz')
#     vol = nib.load(seg_path)
#     imag = vol.get_fdata()
    
#     for i in range(imag.shape[0]):
#         for j in range(imag.shape[1]):
#             for k in range(imag.shape[2]):
#                 if imag[i,j,k]==1.0:
#                     imag[i,j,k]=0
#                 elif imag[i,j,k]==2.0:
#                     imag[i,j,k]=1.0

#     # vol = np.swapaxes(vol, 0, 2)
#     # get image data
#     nifti = nib.Nifti1Image(imag, None)
    
    data_dir = os.path.join('../autodl-tmp/augment_point_traindata',file)
    if not os.path.isdir(data_dir):
        os.makedirs(data_dir)
    image_path = os.path.join('../autodl-tmp/point_train_data/',file,'imaging.nii.gz')
    imaging = nib.load(image_path)
    imaging = imaging.get_fdata()
    w_width = 800
    w_center = 100
    imaging = adjustMethod1(imaging,w_width,w_center)
    imaging_nib = nib.Nifti1Image(imaging,None)
    nib.save(imaging_nib,os.path.join(data_dir,'imaging.nii.gz'))
    # Save segmentation to disk
    # nib.save(nifti, os.path.join(data_dir,'segmentation.nii.gz'))
    # sitk.WriteImage(image_out,os.path.join(data_dir,'segmentation.nii.gz'))
